import torch
import torch.nn as nn
import torch.nn.functional as F


class GCN(nn.Module):
    def __init__(self, in_ft, out_ft, act, bias=True):
        super(GCN, self).__init__()
        self.fc = nn.Linear(in_ft, out_ft, bias=False)
        self.act = nn.PReLU() if act == 'prelu' else act

        if bias:
            self.bias = nn.Parameter(torch.FloatTensor(out_ft))
            self.bias.data.fill_(0.0)
        else:
            self.register_parameter('bias', None)

        for m in self.modules():
            self.weights_init(m)

    @staticmethod
    def weights_init(m):
        if isinstance(m, nn.Linear):
            torch.nn.init.xavier_uniform_(m.weight.data)
            if m.bias is not None:
                m.bias.data.fill_(0.0)

    # Shape of seq: (batch, nodes, features)
    def forward(self, seq, adj, sparse=False):
        seq_fts = self.fc(seq)
        if sparse:
            out = torch.unsqueeze(torch.spmm(adj, torch.squeeze(seq_fts, 0)), 0)
        else:
            out = torch.bmm(adj, seq_fts)
        if self.bias is not None:
            out += self.bias
        out = torch.squeeze(self.act(out), 0)
        return self.act(out)


class HGCN(nn.Module):
    def __init__(self, n_input, hid, semantic_hid, p, act):
        """Dense version of GAT."""
        super(HGCN, self).__init__()
        self.gcn_level_embeddings = []
        self.p = p  # number of meta-Path
        for _ in range(p):
            self.gcn_level_embeddings.append(GCN(n_input, hid, act, bias=True))

        for i, gcn_embedding_path in enumerate(self.gcn_level_embeddings):
            self.add_module('gcn_path_{}'.format(i), gcn_embedding_path)

        self.semantic_level_attention = SemanticAttentionLayer(hid, semantic_hid)

    def forward(self, x, adjs, sparse):
        meta_path_x = []
        for i, adj in enumerate(adjs):
            m_x = self.gcn_level_embeddings[i](x, adj, sparse)
            meta_path_x.append(m_x)

        x = torch.cat([m_x for m_x in meta_path_x], dim=0)

        x = self.semantic_level_attention(x, self.p)
        x = torch.unsqueeze(x, 0)
        return x


class SemanticAttentionLayer(nn.Module):

    def __init__(self, in_features, out_features):
        super(SemanticAttentionLayer, self).__init__()
        self.in_features = in_features
        self.out_features = out_features

        self.W = nn.Parameter(torch.zeros(size=(in_features, out_features)))
        nn.init.xavier_uniform_(self.W.data, gain=1.414)
        self.b = nn.Parameter(torch.zeros(size=(1, out_features)))
        nn.init.xavier_uniform_(self.b.data, gain=1.414)
        self.q = nn.Parameter(torch.zeros(size=(1, out_features)))
        nn.init.xavier_uniform_(self.q.data, gain=1.414)
        self.Tanh = nn.Tanh()

    # input (PN)*F
    def forward(self, x, p):
        h = torch.mm(x, self.W)
        h_prime = self.Tanh(h + self.b.repeat(h.size()[0], 1))
        semantic_attentions = torch.mm(h_prime, torch.t(self.q)).view(p, -1)
        n = semantic_attentions.size()[1]
        semantic_attentions = semantic_attentions.mean(dim=1, keepdim=True)
        semantic_attentions = F.softmax(semantic_attentions, dim=0)
        semantic_attentions = semantic_attentions.view(p, 1, 1)
        semantic_attentions = semantic_attentions.repeat(1, n, self.in_features)
        input_embedding = x.view(p, n, self.in_features)

        h_embedding = torch.mul(input_embedding, semantic_attentions)
        h_embedding = torch.sum(h_embedding, dim=0).squeeze()

        return h_embedding


class AvgReadout(nn.Module):
    def __init__(self):
        super(AvgReadout, self).__init__()

    def forward(self, seq, msk):
        if msk is None:
            return torch.mean(seq, 1)
        else:
            msk = torch.unsqueeze(msk, -1)
            return torch.sum(seq * msk, 1) / torch.sum(msk)


class Discriminator(nn.Module):
    def __init__(self, n_h):
        super(Discriminator, self).__init__()
        self.f_k = nn.Bilinear(n_h, n_h, 1)

        for m in self.modules():
            self.weights_init(m)

    @staticmethod
    def weights_init(m):
        if isinstance(m, nn.Bilinear):
            torch.nn.init.xavier_uniform_(m.weight.data)
            if m.bias is not None:
                m.bias.data.fill_(0.0)

    def forward(self, c, h_pl, h_mi, s_bias1=None, s_bias2=None):
        c_x = torch.unsqueeze(c, 1)
        c_x = c_x.expand_as(h_pl)

        sc_1 = torch.squeeze(self.f_k(h_pl, c_x), 2)
        sc_2 = torch.squeeze(self.f_k(h_mi, c_x), 2)

        if s_bias1 is not None:
            sc_1 += s_bias1
        if s_bias2 is not None:
            sc_2 += s_bias2

        logits = torch.cat((sc_1, sc_2), 1)

        return logits


class LogReg(nn.Module):
    def __init__(self, ft_in, nb_classes):
        super(LogReg, self).__init__()
        self.fc = nn.Linear(ft_in, nb_classes)

        for m in self.modules():
            self.weights_init(m)

    @staticmethod
    def weights_init(m):
        if isinstance(m, nn.Linear):
            torch.nn.init.xavier_uniform_(m.weight.data)
            if m.bias is not None:
                m.bias.data.fill_(0.0)

    def forward(self, seq):
        ret = self.fc(seq)
        return ret
